'''
Parse confusion_matrix to get a numpy ndarray of the matrix
'''

import numpy as np
import os

DIRNAME = os.path.dirname(os.path.abspath(__file__))

def get_conf_mat(label_num, lines_per_row, DATASET=""):
    '''
    input:
        label_num: the number of labels in this dataset
        lines_per_row: the number of lines per row in the confusion matrix
        DATASET: the directory name to find confusion matrix
    return:
        conf_mat: an ndarray of confusion matrix
    '''
    with open(os.path.join(DIRNAME, DATASET+'/confusion_matrix'), 'r', encoding='utf-8') as f:
        def parse_line(line):
            return list(map(int, line.strip().strip('[').strip(']').split()))
        conf_mat =[]
        for i in range(label_num):
            row = []
            for j in range(lines_per_row):
                line = parse_line(f.readline())
                row.extend(line)
            assert len(row) == label_num, f"row {i} is not label_num"
            conf_mat.append(row)
        conf_mat = np.array(conf_mat)
        return conf_mat

if __name__ == '__main__':
    np.set_printoptions(threshold=np.inf)
    print(get_conf_mat())
